package org.kaifamiao.currentlimiting;

import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.time.Duration;
import java.time.Instant;
import java.util.LinkedList;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class SlidingWindowRateLimiterInterceptor implements HandlerInterceptor {

    private final ConcurrentMap<String, LinkedList<Instant>> requestTimestamps = new ConcurrentHashMap<>();
    private final int maxRequestsPerMinute;
    private final Duration timeWindow;

    public SlidingWindowRateLimiterInterceptor(int maxRequestsPerMinute) {
        this.maxRequestsPerMinute = maxRequestsPerMinute;
        this.timeWindow = Duration.ofMinutes(1);
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String clientIp = getClientIP(request);
        LinkedList<Instant> timestamps = requestTimestamps.computeIfAbsent(clientIp, k -> new LinkedList<>());

        synchronized (timestamps) {
            Instant now = Instant.now();
            timestamps.removeIf(timestamp -> Duration.between(timestamp, now).compareTo(timeWindow) > 0);

            if (timestamps.size() < maxRequestsPerMinute) {
                timestamps.add(now);
                return true;
            } else {
                sendTooManyRequestsError(response);
                return false;
            }
        }
    }

    private void sendTooManyRequestsError(HttpServletResponse response) throws IOException {
        response.setContentType("text/plain");
        response.setCharacterEncoding("UTF-8");
        response.setStatus(429);

        try (PrintWriter writer = response.getWriter()) {
            writer.write("请求超过最大次数");
        }
    }

    private String getClientIP(HttpServletRequest request) {
        String ipAddress = request.getHeader("X-Forwarded-For");
        if (ipAddress == null) {
            ipAddress = request.getRemoteAddr();
        }
        return ipAddress;
    }
}
